{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fitting Thomson Scattering Spectra"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[spectral_density]: ../../api/plasmapy.diagnostics.thomson.spectral_density.rst#plasmapy.diagnostics.thomson.spectral_density\n",
    "[lmfit]: https://lmfit.github.io/lmfit-py/\n",
    "\n",
    "Thomson scattering diagnostics record a scattered power spectrum that encodes information about the electron and ion density, temperatures, and flow velocities. This information can be retrieved by fitting the measured spectrum with the theoretical spectral density ([spectral_density]). This notebook demonstrates how to use the [lmfit] package (along with some helpful PlasmaPy functions) to fit 1D Thomson scattering spectra. \n",
    "\n",
    "<img src=\"collective_vs_noncollective_ots.png\">\n",
    "\n",
    "Thomson scattering can be either non-collective (dominated by single electron scattering) or collective (dominated by scattering off of electron plasma waves (EPW) and ion acoustic waves (IAW). In the non-collective regime, the scattering spectrum contains a single peak. However, in the collective regime the spectrum contains separate features caused by the electron and ion populations (corresponding to separate scattering off of EPW and IAW). These features exist on different scales: the EPW feature is dim but covers a wide wavelength range, while the IAW feature is bright but narrow. They also encode partially-degenerate information (e.g., the flow velocities of the electrons and ions respectively). The two features are therefore often recorded on separate spectrometers and are fit separately. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "import astropy.units as u\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from lmfit import Parameters\n",
    "\n",
    "from plasmapy.diagnostics import thomson"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Contents"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "1. [Fitting Collective Thomson Scattering](#Fitting-Collective-Thomson-Scattering)\n",
    "    1. [Fitting the EPW Feature](#Fitting-the-EPW-Feature)\n",
    "    2. [Fitting the IAW Feature](#Fitting-the-IAW-Feature)\n",
    "2. [Fitting Non-Collective Thomson Scattering](#Fitting-Non-Collective-Thomson-Scattering)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fitting Collective Thomson Scattering"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[spectral_density]: ../../api/plasmapy.diagnostics.thomson.spectral_density.rst#plasmapy.diagnostics.thomson.spectral_density\n",
    "\n",
    "To demonstrate the fitting capabilities, we'll first generate some synthetic Thomson data using the [spectral_density] function. This data will be in the collective regime, so we will generate two datasets (using the same plasma parameters and probe geometry) that correspond to the EPW and IAW features. For more details on the spectral density function, see the spectral density function example notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate theoretical spectrum\n",
    "probe_wavelength = 532 * u.nm\n",
    "epw_wavelengths = (\n",
    "    np.linspace(probe_wavelength.value - 30, probe_wavelength.value + 30, num=500)\n",
    "    * u.nm\n",
    ")\n",
    "iaw_wavelengths = (\n",
    "    np.linspace(probe_wavelength.value - 3, probe_wavelength.value + 3, num=500) * u.nm\n",
    ")\n",
    "\n",
    "probe_vec = np.array([1, 0, 0])\n",
    "scattering_angle = np.deg2rad(63)\n",
    "scatter_vec = np.array([np.cos(scattering_angle), np.sin(scattering_angle), 0])\n",
    "\n",
    "n = 2e17 * u.cm**-3\n",
    "ions = [\"H+\", \"C-12 5+\"]\n",
    "T_e = 10 * u.eV\n",
    "T_i = np.array([20, 50]) * u.eV\n",
    "electron_vel = np.array([[0, 0, 0]]) * u.km / u.s\n",
    "ion_vel = np.array([[0, 0, 0], [200, 0, 0]]) * u.km / u.s\n",
    "ifract = [0.3, 0.7]\n",
    "\n",
    "alpha, epw_skw = thomson.spectral_density(\n",
    "    epw_wavelengths,\n",
    "    probe_wavelength,\n",
    "    n,\n",
    "    T_e=T_e,\n",
    "    T_i=T_i,\n",
    "    ions=ions,\n",
    "    ifract=ifract,\n",
    "    electron_vel=electron_vel,\n",
    "    ion_vel=ion_vel,\n",
    "    probe_vec=probe_vec,\n",
    "    scatter_vec=scatter_vec,\n",
    ")\n",
    "\n",
    "alpha, iaw_skw = thomson.spectral_density(\n",
    "    iaw_wavelengths,\n",
    "    probe_wavelength,\n",
    "    n,\n",
    "    T_e=T_e,\n",
    "    T_i=T_i,\n",
    "    ions=ions,\n",
    "    ifract=ifract,\n",
    "    electron_vel=electron_vel,\n",
    "    ion_vel=ion_vel,\n",
    "    probe_vec=probe_vec,\n",
    "    scatter_vec=scatter_vec,\n",
    ")\n",
    "\n",
    "# PLOTTING\n",
    "fig, ax = plt.subplots(ncols=2, figsize=(12, 4))\n",
    "fig.subplots_adjust(wspace=0.2)\n",
    "\n",
    "for a in ax:\n",
    "    a.set_xlabel(\"Wavelength (nm)\")\n",
    "    a.set_ylabel(\"Skw\")\n",
    "    a.axvline(x=probe_wavelength.value, color=\"red\", label=\"Probe wavelength\")\n",
    "\n",
    "ax[0].set_xlim(520, 545)\n",
    "ax[0].set_ylim(0, 3e-13)\n",
    "ax[0].set_title(\"Electron Plasma Wave\")\n",
    "ax[0].plot(epw_wavelengths.value, epw_skw)\n",
    "ax[0].legend()\n",
    "\n",
    "ax[1].set_xlim(530, 534)\n",
    "ax[1].set_title(\"Ion Acoustic Wave\")\n",
    "ax[1].plot(iaw_wavelengths.value, iaw_skw)\n",
    "ax[1].legend();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that these plots are showing the same spectral distribution, just over different wavelength ranges: the large peak in the center of the EPW spectrum is the IAW spectrum. Next we'll add some noise to the spectra to simulate an experimental measurement."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "epw_skw *= 1 + np.random.normal(loc=0, scale=0.1, size=epw_wavelengths.size)\n",
    "\n",
    "iaw_skw *= 1 + np.random.normal(loc=0, scale=0.1, size=iaw_wavelengths.size)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[np.nan]: https://numpy.org/doc/stable/reference/constants.html#numpy.nan\n",
    "During experiments, the IAW feature is typically blocked on the EPW detector using a notch filter to prevent it from saturating the measurement. We'll mimic this by setting the center of the EPW spectrum to [np.nan]. The fitting algorithm applied later will recognize these values and not include them in the fit.\n",
    "\n",
    "The more of the EPW spectrum we exclude, the less sensitive the fit will become. So, we want to block as little as possible while still obscuring the IAW feature."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "notch_range = (531, 533)\n",
    "x0 = np.argmin(np.abs(epw_wavelengths.value - notch_range[0]))\n",
    "x1 = np.argmin(np.abs(epw_wavelengths.value - notch_range[1]))\n",
    "epw_skw[x0:x1] = np.nan"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[np.nanmax()]: https://numpy.org/doc/stable/reference/generated/numpy.nanmax.html\n",
    "\n",
    "Another option is to delete the missing data from both the data vector and the corresponding wavelengths vector. \n",
    "```\n",
    "epw_skw = np.delete(epw_skw, slice(x0,x1,None))\n",
    "epw_wavelengths = np.delete(epw_wavelengths, slice(x0,x1,None))\n",
    "```\n",
    "This does not look as nice on plots (since it leaves a line between datapoints at ``x0`` and ``x1``) but is necessary in some cases (e.g. when also using an instrument function). \n",
    "\n",
    "Finally, we need to get rid of the units and normalize the data on each detector to its maximum value (this is a requirement of the fitting algorithm). We're using [np.nanmax()] here to ignore the NaN values in `epw_skw`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "epw_skw = epw_skw.value\n",
    "epw_skw *= 1 / np.nanmax(epw_skw)\n",
    "iaw_skw = iaw_skw.value\n",
    "iaw_skw *= 1 / np.nanmax(iaw_skw)\n",
    "\n",
    "# Plot again\n",
    "fig, ax = plt.subplots(ncols=2, figsize=(12, 4))\n",
    "fig.subplots_adjust(wspace=0.2)\n",
    "\n",
    "for a in ax:\n",
    "    a.set_xlabel(\"Wavelength (nm)\")\n",
    "    a.set_ylabel(\"Skw\")\n",
    "    a.axvline(x=probe_wavelength.value, color=\"red\", label=\"Probe wavelength\")\n",
    "\n",
    "ax[0].set_xlim(520, 545)\n",
    "ax[0].set_title(\"Electron Plasma Wave\")\n",
    "ax[0].plot(epw_wavelengths.value, epw_skw)\n",
    "ax[0].legend()\n",
    "\n",
    "ax[1].set_xlim(531, 533)\n",
    "ax[1].set_title(\"Ion Acoustic Wave\")\n",
    "ax[1].plot(iaw_wavelengths.value, iaw_skw)\n",
    "ax[1].legend();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll start by fitting the EPW feature, then move on to the IAW feature. This is typically the process followed when analyzing experimental data, since the EPW feature depends on fewer parameters."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Fitting the EPW Feature"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[spectral_density]: ../../api/plasmapy.diagnostics.thomson.spectral_density.rst#plasmapy.diagnostics.thomson.spectral_density\n",
    "[lmfit]: https://lmfit.github.io/lmfit-py/\n",
    "[lmfit.Parameter]: https://lmfit.github.io/lmfit-py/parameters.html#lmfit.parameter.Parameter\n",
    "[lmfit.Parameters]: https://lmfit.github.io/lmfit-py/parameters.html#lmfit.parameter.Parameters\n",
    "[lmfit.Model]:https://lmfit.github.io/lmfit-py/model.html#lmfit.model.Model\n",
    "[spectral_density_model]: ../../api/plasmapy.diagnostics.thomson.spectral_density_model.rst#plasmapy.diagnostics.thomson.spectral_density_model\n",
    "\n",
    "In order to fit this data in [lmfit] we need to create an [lmfit.Model] object for this problem. PlasmaPy includes such a model function for fitting Thomson scattering spectra, [spectral_density_model]. This model function is initialized by providing the parameters of the fit in the form of an [lmfit.Parameters] object.\n",
    "\n",
    "A [lmfit.Parameters] object is an ordered dictionary of [lmfit.Parameter] objects. Each [lmfit.Parameter] has a number of elements, the most important of which for our purposes are\n",
    "- \"name\" (str) -> The name of the parameter (and the key in `Parameters` dictionary)\n",
    "- \"value\" (float) -> The initial value of the parameter.\n",
    "- \"vary\" (boolean) -> Whether or not the parameter will be varied during fitting.\n",
    "- \"min\", \"max\" (float) -> The minimum and maximum bounds for the parameter during fitting.\n",
    "\n",
    "Since [lmfit.Parameter] objects can only be scalars, arrays of multiple quantities must be broken apart into separate `Parameter` objects. To do so, this fitting routine adopts the following convention:\n",
    "\n",
    "`T_e = [1,2] -> \"T_e_0\"=1, \"T_e_1\" = 2`\n",
    "\n",
    "Specifying large arrays (like velocity vectors for multiple ion species) is clearly tedious in this format, and specifying non-numeric values (such as `ions`) is impossible. Therefore, this routine also takes a `settings` dictionary as a way to pass non-varying input to the [spectral_density_model] function.\n",
    "\n",
    "A list of required and optional parameters and settings is provided in the docstring for the [spectral_density_model] function. For example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "help(thomson.spectral_density_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[lmfit.Parameters]: https://lmfit.github.io/lmfit-py/parameters.html#lmfit.parameter.Parameters\n",
    "[spectral_density]: ../../api/plasmapy.diagnostics.thomson.spectral_density.rst#plasmapy.diagnostics.thomson.spectral_density\n",
    "\n",
    "We will now create the [lmfit.Parameters] object and settings dictionary using the values defined earlier when creating the sample data. We will choose to vary both the density and electron temperature, and will intentionally set incorrect initial values for both parameters. Note that, even though only one electron population is included, we must still name the temperature variable `Te_0` in accordance with the convention defined above.\n",
    "\n",
    "Note that the EPW spectrum is effectively independent of the ion parameters: only `n`, `Te`, and `electron_vel` will affect this fit. However, since the ion parameters are required arguments for [spectral_density], we still need to provide values for them. We will therefore set these parameters as fixed (`vary=False`) with approximate values (in this case, intentionally poor estimates have been chosen to emphasize that they do not affect the EPW fit). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "params = Parameters()\n",
    "params.add(\n",
    "    \"n\", value=4e17 * 1e6, vary=True, min=5e16 * 1e6, max=1e18 * 1e6\n",
    ")  # Converting cm^-3 to m^-3\n",
    "params.add(\"T_e_0\", value=5, vary=True, min=0.5, max=25)\n",
    "params.add(\"T_i_0\", value=5, vary=False)\n",
    "params.add(\"T_i_1\", value=10, vary=False)\n",
    "params.add(\"ifract_0\", value=0.8, vary=False)\n",
    "params.add(\"ifract_1\", value=0.2, vary=False)\n",
    "params.add(\"ion_speed_0\", value=0, vary=False)\n",
    "params.add(\"ion_speed_1\", value=0, vary=False)\n",
    "\n",
    "settings = {}\n",
    "settings[\"probe_wavelength\"] = probe_wavelength.to(u.m).value\n",
    "settings[\"probe_vec\"] = probe_vec\n",
    "settings[\"scatter_vec\"] = scatter_vec\n",
    "settings[\"ions\"] = ions\n",
    "settings[\"ion_vdir\"] = np.array([[1, 0, 0], [1, 0, 0]])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[spectral_density]: ../../api/plasmapy.diagnostics.thomson.spectral_density.rst#plasmapy.diagnostics.thomson.spectral_density\n",
    "[lmfit]: https://lmfit.github.io/lmfit-py/\n",
    "[spectral_density_model]: ../../api/plasmapy.diagnostics.thomson.spectral_density_model.rst#plasmapy.diagnostics.thomson.spectral_density_model\n",
    "[constraint]: https://lmfit.github.io/lmfit-py/constraints.html\n",
    "[lmfit.Parameter]: https://lmfit.github.io/lmfit-py/parameters.html#lmfit.parameter.Parameter\n",
    "[lmfit.Model]:https://lmfit.github.io/lmfit-py/model.html#lmfit.model.Model\n",
    "\n",
    "[lmfit] allows the value of a parameter to be fixed using a [constraint] using the `expr` keyword in [lmfit.Parameter] as shown in the definition of `ifract_1` above. In this case, `ifract_1` is not actually a free parameter, since its value is fixed by the fact that `ifract_0 + ifract_1 = 1.0`. This constraint is made explicit here, but the [spectral_density_model] function will automatically enforce this constraint for `efract` and `ifract` variables.\n",
    "\n",
    "Just as in the [spectral_density] function, some parameters are required while others are optional. For example, density (`n`) is required but `ion_velocity` is optional (and will be assumed to be zero if not provided). The list of required and optional parameters is identical to the required and optional arguments of [spectral_density].\n",
    "\n",
    "We can now use these objects to initialize an [lmfit.Model] object based on the [spectral_density] function using the [spectral_density_model] function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "epw_model = thomson.spectral_density_model(\n",
    "    epw_wavelengths.to(u.m).value, settings, params\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[model.fit()]: https://lmfit.github.io/lmfit-py/model.html#lmfit.model.Model.fit\n",
    "[minimizer options]: https://lmfit.github.io/lmfit-py/fitting.html#lmfit.minimizer.minimize\n",
    "[np.ndarray]:https://numpy.org/doc/stable/reference/generated/numpy.ndarray.html\n",
    "\n",
    "With the model created, the fit can be easily performed using the [model.fit()] method. This method takes several keyword options that are worth mentioning: \n",
    "\n",
    "- \"method\" -> A string that defines the fitting method, from the list of [minimizer options]. A good choice for fitting Thomson spectra is `differential_evolution`.\n",
    "\n",
    "- \"max_nfev\" -> The maximum number of iterations allowed.\n",
    "\n",
    "In addition, of course we also need to include the data to be fit (`epw_skw`), the independent variable (`wavelengths`) and the parameter object. It is important to note that the data to be fit should be a [np.ndarray] (unit-less) and normalized."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fit_kws = {}\n",
    "epw_result = epw_model.fit(\n",
    "    epw_skw,\n",
    "    params=params,\n",
    "    wavelengths=epw_wavelengths.to(u.m).value,\n",
    "    method=\"differential_evolution\",\n",
    "    fit_kws=fit_kws,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[model.fit()]: https://lmfit.github.io/lmfit-py/model.html#lmfit.model.Model.fit\n",
    "[scipy.optimize]: https://docs.scipy.org/doc/scipy/reference/optimize.html\n",
    "[differential_evolution]:https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.differential_evolution.html\n",
    "[lmfit.ModelResult]:https://lmfit.github.io/lmfit-py/model.html#lmfit.model.ModelResult\n",
    "\n",
    "In the [model.fit()], the `fit_kws` keyword can be set to a dictionary of keyword arguments that will be passed to the underlying [scipy.optimize] optimization function, and can be used to refine the fit. For details, see the SciPy documentation for the [differential_evolution] algorithm. \n",
    "\n",
    "The return from this function is a [lmfit.ModelResult] object, which is a convenient container holding lots of information! To start with, we can see the best-fit parameters, number of iterations, the chiSquared goodness-of-fit metric, and plot the best-fit curve."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Print some of the results compared to the true values\n",
    "answers = {\"n\": 2e17, \"T_e_0\": 10}\n",
    "for key, ans in answers.items():\n",
    "    print(f\"{key}: {epw_result.best_values[key]:.1e} (true value {ans:.1e})\")\n",
    "\n",
    "print(f\"Number of fit iterations:{epw_result.nfev}\")\n",
    "print(f\"Reduced Chisquared:{epw_result.redchi:.4f}\")\n",
    "\n",
    "# Extract the best fit curve by evaluating the model at the final parameters\n",
    "n_fit = epw_result.values[\"n\"]\n",
    "Te_0_fit = epw_result.values[\"T_e_0\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that the best_fit curve skips the NaN values in the data, so the array is shorter than `epw_skw`. In order to plot them against the same wavelengths, we need to create an array of indices where `epw_skw` is not NaN."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "nbsphinx-thumbnail": {
     "tooltip": "Thomson Scattering"
    }
   },
   "outputs": [],
   "source": [
    "# Extract the best fit curve\n",
    "best_fit_skw = epw_result.best_fit\n",
    "\n",
    "# Get all the non-nan indices (the best_fit_skw just omits these values)\n",
    "not_nan = np.argwhere(np.logical_not(np.isnan(epw_skw)))\n",
    "\n",
    "# Plot\n",
    "fig, ax = plt.subplots(ncols=1, figsize=(8, 8))\n",
    "ax.set_xlabel(\"Wavelength (nm)\")\n",
    "ax.set_ylabel(\"Skw\")\n",
    "ax.axvline(x=probe_wavelength.value, color=\"red\", label=\"Probe wavelength\")\n",
    "\n",
    "ax.set_xlim(520, 545)\n",
    "\n",
    "ax.plot(epw_wavelengths.value, epw_skw, label=\"Data\")\n",
    "ax.plot(epw_wavelengths.value[not_nan], best_fit_skw, label=\"Best-fit\")\n",
    "ax.legend(loc=\"upper right\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The resulting fit is very good, even though many of the ion parameters are still pretty far from their actual values!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fitting the IAW Feature"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[lmfit.Parameters]: https://lmfit.github.io/lmfit-py/parameters.html#lmfit.parameter.Parameters\n",
    "\n",
    "We will now follow the same steps to fit the IAW feature. We'll start by setting up a new [lmfit.Parameters] object, this time using the best-fit values from the EPW fit as fixed values for those parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "settings = {}\n",
    "settings[\"probe_wavelength\"] = probe_wavelength.to(u.m).value\n",
    "settings[\"probe_vec\"] = probe_vec\n",
    "settings[\"scatter_vec\"] = scatter_vec\n",
    "settings[\"ions\"] = ions\n",
    "settings[\"ion_vdir\"] = np.array([[1, 0, 0], [1, 0, 0]])\n",
    "\n",
    "params = Parameters()\n",
    "params.add(\"n\", value=n_fit, vary=False)\n",
    "params.add(\"T_e_0\", value=Te_0_fit, vary=False)\n",
    "params.add(\"T_i_0\", value=10, vary=True, min=5, max=60)\n",
    "params.add(\"T_i_1\", value=10, vary=True, min=5, max=60)\n",
    "params.add(\"ifract_0\", value=0.5, vary=True, min=0.2, max=0.8)\n",
    "params.add(\"ifract_1\", value=0.5, vary=True, min=0.2, max=0.8, expr=\"1.0 - ifract_0\")\n",
    "params.add(\"ion_speed_0\", value=0, vary=False)\n",
    "params.add(\"ion_speed_1\", value=0, vary=True, min=0, max=1e6)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we will run the fit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "iaw_model = thomson.spectral_density_model(\n",
    "    iaw_wavelengths.to(u.m).value, settings, params\n",
    ")\n",
    "\n",
    "iaw_result = iaw_model.fit(\n",
    "    iaw_skw,\n",
    "    params=params,\n",
    "    wavelengths=iaw_wavelengths.to(u.m).value,\n",
    "    method=\"differential_evolution\",\n",
    ")\n",
    "\n",
    "# Print some of the results compared to the true values\n",
    "answers = {\n",
    "    \"T_i_0\": 20,\n",
    "    \"T_i_1\": 50,\n",
    "    \"ifract_0\": 0.3,\n",
    "    \"ifract_1\": 0.7,\n",
    "    \"ion_speed_1\": 2e5,\n",
    "}\n",
    "for key, ans in answers.items():\n",
    "    print(f\"{key}: {iaw_result.best_values[key]:.1f} (true value {ans:.1f})\")\n",
    "\n",
    "print(f\"Number of fit iterations:{iaw_result.nfev:.1f}\")\n",
    "print(f\"Reduced Chisquared:{iaw_result.redchi:.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And plot the results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Extract the best fit curve\n",
    "best_fit_skw = iaw_result.best_fit\n",
    "\n",
    "# Plot\n",
    "fig, ax = plt.subplots(ncols=1, figsize=(8, 8))\n",
    "ax.set_xlabel(\"Wavelength (nm)\")\n",
    "ax.set_ylabel(\"Skw\")\n",
    "ax.axvline(x=probe_wavelength.value, color=\"red\", label=\"Probe wavelength\")\n",
    "\n",
    "ax.set_xlim(531, 533)\n",
    "\n",
    "ax.plot(iaw_wavelengths.value, iaw_skw, label=\"Data\")\n",
    "ax.plot(iaw_wavelengths.value, best_fit_skw, label=\"Best-fit\")\n",
    "ax.legend(loc=\"upper right\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "At this point, the best-fit parameters from the IAW fit could be used to further refine the EPW fit, and so on iteratively until both fits become stable."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fitting Non-Collective Thomson Scattering"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The non-collective Thomson scattering spectrum depends only on the electron density and temperature. In this regime the spectrum is less featured and, consequently, fits will produce larger errors. Otherwise, the fitting procedure is the same. \n",
    "\n",
    "To illustrate fitting in this regime, we'll start by generating some test data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nc_wavelengths = (\n",
    "    np.linspace(probe_wavelength.value - 50, probe_wavelength.value + 50, num=1000)\n",
    "    * u.nm\n",
    ")\n",
    "\n",
    "n = 5e15 * u.cm**-3\n",
    "T_e = 5 * u.eV\n",
    "T_i = np.array([1]) * u.eV\n",
    "ions = [\"H+\"]\n",
    "\n",
    "alpha, Skw = thomson.spectral_density(\n",
    "    nc_wavelengths,\n",
    "    probe_wavelength,\n",
    "    n,\n",
    "    T_e=T_e,\n",
    "    T_i=T_i,\n",
    "    ions=ions,\n",
    "    probe_vec=probe_vec,\n",
    "    scatter_vec=scatter_vec,\n",
    ")\n",
    "\n",
    "# Normalize and add noise\n",
    "nc_skw = Skw.value\n",
    "nc_skw *= 1 / np.nanmax(nc_skw)\n",
    "nc_theory = np.copy(nc_skw)\n",
    "nc_skw *= 1 + np.random.normal(loc=0, scale=0.1, size=nc_wavelengths.size)\n",
    "\n",
    "\n",
    "# Plot again\n",
    "fig, ax = plt.subplots(ncols=2, figsize=(12, 4))\n",
    "fig.subplots_adjust(wspace=0.2)\n",
    "\n",
    "for a in ax:\n",
    "    a.set_xlabel(\"Wavelength (nm)\")\n",
    "    a.set_ylabel(\"Skw\")\n",
    "    a.axvline(x=probe_wavelength.value, color=\"red\", label=\"Probe wavelength\")\n",
    "\n",
    "ax[0].set_xlim(520, 545)\n",
    "ax[0].set_title(\"Theoretical signal\")\n",
    "ax[0].plot(nc_wavelengths.value, nc_theory)\n",
    "ax[0].legend()\n",
    "\n",
    "ax[1].set_xlim(520, 545)\n",
    "ax[1].set_title(\"With noise\")\n",
    "ax[1].plot(nc_wavelengths.value, nc_skw)\n",
    "ax[1].legend();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[lmfit.Parameters]: https://lmfit.github.io/lmfit-py/parameters.html#lmfit.parameter.Parameters\n",
    "[spectral_density]: ../../api/plasmapy.diagnostics.thomson.spectral_density.rst#plasmapy.diagnostics.thomson.spectral_density\n",
    "Then we will setup the settings dictionary and [lmfit.Parameters] object and run the fit.\n",
    "\n",
    "\n",
    "Note that, in the non-collective regime, the [spectral_density] function is effectively independent of density. Just as in the [collective-regime EPW fitting](#Fitting-the-EPW-Feature) section, we still need to provide a value for `n` because it is a required argument for [spectral_density], but this parameter should be fixed and its value can be a rough approximation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "settings = {}\n",
    "settings[\"probe_wavelength\"] = probe_wavelength.to(u.m).value\n",
    "settings[\"probe_vec\"] = probe_vec\n",
    "settings[\"scatter_vec\"] = scatter_vec\n",
    "settings[\"ions\"] = ions\n",
    "\n",
    "params = Parameters()\n",
    "params.add(\"n\", value=1e15 * 1e6, vary=False)  # Converting cm^-3 to m^-3\n",
    "params.add(\"T_e_0\", value=1, vary=True, min=0.1, max=20)\n",
    "params.add(\"T_i_0\", value=1, vary=False)\n",
    "\n",
    "nc_model = thomson.spectral_density_model(\n",
    "    nc_wavelengths.to(u.m).value, settings, params\n",
    ")\n",
    "\n",
    "nc_result = nc_model.fit(\n",
    "    nc_skw,\n",
    "    params,\n",
    "    wavelengths=nc_wavelengths.to(u.m).value,\n",
    "    method=\"differential_evolution\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "best_fit_skw = nc_result.best_fit\n",
    "\n",
    "print(f\"T_e_0: {nc_result.best_values['T_e_0']:.1f} (true value 5)\")\n",
    "print(f\"Number of fit iterations:{nc_result.nfev}\")\n",
    "print(f\"Reduced Chisquared:{nc_result.redchi:.4f}\")\n",
    "\n",
    "# Plot\n",
    "fig, ax = plt.subplots(ncols=1, figsize=(8, 8))\n",
    "ax.set_xlabel(\"Wavelength (nm)\")\n",
    "ax.set_ylabel(\"Skw\")\n",
    "ax.axvline(x=probe_wavelength.value, color=\"red\", label=\"Probe wavelength\")\n",
    "\n",
    "ax.set_xlim(520, 545)\n",
    "\n",
    "ax.plot(nc_wavelengths.value, nc_skw, label=\"Data\")\n",
    "ax.plot(nc_wavelengths.value, best_fit_skw, label=\"Best-fit\")\n",
    "ax.legend(loc=\"upper right\");"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
